import pandas as pd
import pandas_datareader as pdr
import matplotlib as mpl
import matplotlib.pyplot as plt
from sklearn.svm import SVR
from scipy.spatial.distance import euclidean
from fastdtw import fastdtw


def get_sp500_data():
    """
    标普500指数
    日期，开盘，最高，最低，收盘，成交量，调整收盘价
    日变化：收盘-开盘
    隔夜变化：开盘 - 前收盘
    :return:
    """
    df = pd.read_excel('./data/sp500.xlsx')
    df.sort_values('Date', inplace=True)
    df.reset_index(drop=True, inplace=True)
    return df


def get_yahoo_date(code, start_date, end_date):
    return pdr.get_data_yahoo(code, start_date, end_date)


def sp500_plot(df):
    """
    sp500收盘价画图
    :param df:
    :return:
    """
    # matplotlib中文显示方块
    mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
    mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题

    # df_sp500_plot = df[['Date', 'Close']]
    # df_sp500_plot.index = df_sp500_plot['Date']
    # df_sp500_plot['Close'].plot()
    df['Close'].plot()

    plt.title('标普500指数', fontsize=15)

    plt.show()


def svm_model(df, df_index):
    x_train = df[: -1000]
    y_train = df['Close'].shift(-1)[: -1000]
    x_test = sp20[-1000:]
    y_test = df['Close'].shift(-1)[-1000: ]

    clf = SVR(kernel='linear')
    model = clf.fit(x_train, y_train)
    y_pred = model.predict(x_test)

    tf = pd.DataFrame(list(zip(y_test, y_pred)), columns=['Actual_Next_Close', 'Preddicted_Next_Close'], index=y_test.index)
    # 增加当日收盘和下一交易日开盘
    open = df_index.loc[tf.index]['Open'].shift(-1)
    close = df_index.loc[tf.index]['Close']
    # tf['Next_Open'] = pd.DataFrame(open, index=tf.index)
    # tf['Close'] = pd.DataFrame(close, index=tf.index)
    tf = pd.concat([tf, open, close], axis=1)
    tf.columns = ['Actual_Next_Close', 'Preddicted_Next_Close', 'Next_Open', 'Close']

    # 增加信号
    # 交易信号: 修改交易信号(增加预测值比开盘价高的程度)
    tf['Signal'] = tf.apply(lambda x: 1 if x['Preddicted_Next_Close'] > x['Next_Open'] + 1 else 0, axis=1)
    # 交易信号发出后，计算交易日的日内交易收益率
    tf['Benifit_Rate'] = tf.apply(lambda x: (x['Actual_Next_Close'] - x['Next_Open']) / x['Next_Open']
                                if x['Signal'] == 1 else 0, axis=1)
    # 交易信号发出后，计算交易日的日内交易绝对收益
    tf['Benifit'] = tf.apply(lambda x: x['Actual_Next_Close'] - x['Next_Open'] if x['Signal'] == 1 else 0, axis=1)

    print(tf.head())
    # 绝对收益
    print(tf['Benifit'].sum())


def dtw_dist(ser1, ser2):
    """
    返回两个序列之间的距离

    DTW算法：
    DTW通过把时间序列进行延伸和缩短，来计算两个时间序列性之间的相似性：
    两个时间序列之间，存在一些相似的店，DTW使用所有这些相似点之间的距离和，称之为整路径距离来衡量两个序列之间的相似性
    https://www.cnblogs.com/kemaswill/archive/2013/04/18/3028610.html
    https://www.cnblogs.com/kemaswill/archive/2013/04/18/3029078.html
    :param ser1:
    :param ser2:
    :return:
    """
    distance, path = fastdtw(ser1, ser2, dist=euclidean)
    return distance


def dtw_model(df):
    """
    动态时间规整模型
    :param df:
    :return:
    """
    tseries = []
    # 每个周期长度为5个交易日
    tlen = 5
    pct = df['Close'].pct_change().values

    for i in range(tlen, len(df), tlen):
        tseries.append((pct[i - tlen + 1: i], pct[i]))

    # 计算各序列间的距离
    dist_pairs = []
    for i in range(len(tseries)):
        for j in range(len(tseries)):
            dist = dtw_dist(ser1=tseries[i][0], ser2=tseries[j][0])
            dist_pairs.append((i, j, dist, tseries[i][1], tseries[j][1]))

    dist_frame = pd.DataFrame(dist_pairs, columns=['A', 'B', 'Dist', 'A Ret', 'B Ret'])
    # 去掉相同列并排序
    sf = dist_frame[dist_frame['Dist'] > 0].sort_values(['A', 'B']).reset_index(drop=True)
    sfe = sf[sf['A'] < sf['B']]
    winf = sfe[(sfe['Dist'] < 1)&(sfe['A Ret']) > 0]
    print(winf)


# # 获取数据
# df_sp500 = get_sp500_data()
# # 收盘价画图
# sp500_plot(df_sp500)

# # 获取数据
start_date = pd.to_datetime('2000-01-01')
end_date = pd.to_datetime('2016-01-01')
df_sp500 = get_yahoo_date(code='SPY', start_date=start_date, end_date=end_date)
# 收盘价画图
# sp500_plot(df_sp500)

# 获取收盘价，以及前n天的收盘价，用来预测下一个交易日的收盘价
sp20 = df_sp500[['Close']]
for i in range(1, 21):
    sp20['Close_{}_Days_Ago'.format(i)] = sp20['Close'].shift(i)

sp20 = sp20.dropna(how='any')

#  建模
# svm
# svm_model(df=sp20, df_index=df_sp500)
# 动态时间规整（dynamic time warping）
dtw_model(df=sp20)
